import sys
import numpy as np
import torch
import torch.nn as nn
import ctypes

sys.path.append("../utils")

from utils import OpRunner, compare, verify_result
import time


def main():
    func = OpRunner("./build/libop.so", 7, pint_time=True)
    func.op.run.argtypes = None

    minvalue = 1
    maxvalue = 10
    shape = [9, 80, 3]
    keepdim = False
    epsilon = 1e-12
    for dtype in [np.float16, np.float32]:
    # for dtype in [np.float16]:
        for p in [0.0, 1.0, 2.0, 3.0, float("inf"), float("-inf")]:
        # for p in [1.0]:
            for dim in [-1, 0, 1, 2]:
            # for dim in [-1]:
                # input_x = np.arange(8).reshape(shape).astype(dtype)
                input_x = np.random.uniform(minvalue, maxvalue, shape).astype(dtype)
                output = np.ones(shape, dtype=dtype)
                result = torch.norm(
                    torch.tensor(input_x, dtype=torch.float),
                    p=p,
                    dim=dim if dim >= 0 else None,
                    keepdim=keepdim,
                )
                result = result.numpy()
                golden = result.astype(dtype)

                func(
                    [
                        input_x,
                        ctypes.c_double(p),
                        ctypes.c_int64(dim),
                        ctypes.c_bool(keepdim),
                        ctypes.c_double(epsilon),
                        output,
                    ],
                    [-1],
                )

                if dim < 0:
                    # print(np.transpose(input_x, (2,1,0)))
                    # print(input_x)
                    if verify_result(output[0,0,0], golden):
                        print(dtype, "p =", p, "dim = ", dim, "成功")
                    else:
                        print("output", output[0,0,0])
                        print("golden", golden)
                        print(dtype, "p =", p, "dim = ", dim, "失败")
                    print("----------------------------")
                else:
                    golden = golden.reshape((golden.size))
                    output = output.reshape((output.size))
                    output = output[: golden.size]
                    if verify_result(output, golden):
                        print(dtype, "p =", p, "dim = ", dim, "成功")
                    else:
                        # print("output", output)
                        # print("golden", golden)
                        print(dtype, "p =", p, "dim = ", dim, "失败")
                    print("----------------------------")


if __name__ == "__main__":
    main()
